from Cryptodome.Cipher import AES
from Cryptodome import Random
import os
import chardet  # 获取编码


def aes(data, key):
    iv = Random.new().read(AES.block_size)
    # 使用key和iv初始化AES对象, 使用MODE_CFB模式
    mycipher = AES.new(key, AES.MODE_CFB, iv)
    # 加密的明文长度必须为16的倍数，如果长度不为16的倍数，则需要补足为16的倍数
    # 将iv（密钥向量）加到加密的密文开头，一起传输
    ciphertext = iv + mycipher.encrypt(data.encode())
    return ciphertext


def split(filename, size):
    code = file_code(filename)
    fp = open(filename, 'r', encoding=code)
    i = 0
    n = 0
    temp = open(filename[:-4] + '_part' + str(i) + '.txt', 'w', encoding=code)
    buf = fp.read(1024)
    while (True):
        temp.write(buf)
        buf = fp.read(1024)
        if (buf == ''):
            print(filename[:-4] + '_part' + str(i) + '.txt' + ';')
            temp.close()
            fp.close()
            return i
        n += 1
        if (n == size):
            n = 0
            print(filename[:-4] + '_part' + str(i) + '.txt' + ';')
            i += 1
            temp.close()
            temp = open(filename[:-4] + '_part' + str(i) + '.txt', 'w', encoding=code)


def decrypt(ciphertext, key):
    # 解密的话要用key和iv生成新的AES对象
    mydecrypt = AES.new(key, AES.MODE_CFB, ciphertext[:16])
    # 使用新生成的AES对象，将加密的密文解密
    decrypttext = mydecrypt.decrypt(ciphertext[16:])
    # print(decrypttext)
    return decrypttext.decode()


def file_code(path_file__):
    with open(path_file__, 'rb') as a:  # 获取文件编码
        check = a.read(1)
        code = chardet.detect(check).get('encoding')
        if code != 'utf-8':
            code = 'gbk'
    return code


def aes_files(path_file__, file_ase, key):
    print('\t开始加密文件：' + path_file__)
    code = file_code(path_file__)
    try:
        with open(path_file__, 'r', encoding=code) as file:
            filed = file.read()
            aes_file = aes(filed, key)
            wFile(file_ase, aes_file, 'aes')
    except UnicodeDecodeError:
        print('文件编码错误' + path_file__)
        time.sleep(5)
        return


def aes_file(path, key):
    # key = b'this is a 16 key'
    try:
        file_ase = path + '//file_ase'
        os.mkdir(file_ase)
    except FileExistsError:
        pass
    try:
        for f in os.walk(path):
            print(f)
            if f == None:
                continue
            for i in f[2]:
                if i[-3:] != 'txt':
                    continue
                file_nums = None
                file_ase = path + '\\file_ase\\' + i
                path_file = path + '\\' + i
                file_size = os.path.getsize(path_file)
                if file_size > 5120000:
                    print(path_file + ' ( ' + str(round(file_size / 1024 / 1024, 2)) + 'MB) 文件过大，需进行切割。')
                    file_nums = split(path_file, 128000)
                    try:
                        os.rename(path_file, path_file + '__')
                    except Exception as e:
                        print(e)
                        print('rename file fail\r\n')
                    else:
                        print(path_file + '\trename file success \t' + path_file + '__\r\n')
                print('开始加密文件：' + path_file)
                if file_nums is None:
                    aes_files(path_file, file_ase, key)
                else:
                    num = 0
                    while (file_nums + 1 > num):
                        path_file__ = path_file[:-4] + '_part' + str(num) + '.txt'
                        file_ase__ = file_ase[:-4] + '_part' + str(num) + '.txt'
                        aes_files(path_file__, file_ase__, key)
                        num += 1
                file_nums = None
    except FileNotFoundError:
        if i:
            print('文件加密完成！！！')
        else:
            print('系统未找到指定文件')


def wFile(file_ase, data, mode):
    if mode == 'aes':
        path_file_aes = file_ase[:-4] + '_加密' + '.txt'
        print('成功加密文件：' + path_file_aes)
        print('============================================\n')
        with open(path_file_aes, 'ab+') as f:
            f.write(data)
    else:
        path_file_aes = file_ase[:-7] + '.txt'
        print('成功解密文件：' + path_file_aes)
        print('============================================\n')
        with open(path_file_aes, 'a', encoding='utf-8') as f:
            f.write(data)


def decrypt_file(decrypt_path, key):
    # key = b'this is a 16 key'
    try:
        file_decrypt = decrypt_path + '//file_decrypt'
        os.mkdir(file_decrypt)
    except FileExistsError:
        pass
    try:
        for f in os.walk(decrypt_path):
            for i in f[2]:
                if i[-3:] != 'txt':
                    continue
                file_ase = decrypt_path + '\\file_decrypt\\' + i
                path_file = decrypt_path + '\\' + i
                print('开始解密文件：' + path_file)
                with open(path_file, 'rb') as file:
                    filed = file.read()
                    decrypt_file = decrypt(filed, key)
                    wFile(file_ase, decrypt_file, 'decrypt')
    except FileNotFoundError:
        pass


def aes_key(key):
    if len(key) < 16:
        keyl = '0'.encode() * (16 - len(key))
        key += keyl
    elif len(key) > 16:
        key = key[-16:]
    return key


if __name__ == '__main__':
    import sys
    import time
    from 办公自动化.版权 import tmmeT

    # if tmmeT().ToolC('2021-06-04') != 200:
    #     time.sleep(3)
    #     sys.exit()
    if tmmeT().ToolT('2021-06-04') != 200:
        time.sleep(3)
        sys.exit()
    q = 1
    try:
        while (q != 'q'):
            # key = b'this is a 16 key'
            flag = input('请输入要进行的操作 1、加密\t2、解密\n')
            key = input('请输入密码：').encode()
            key = aes_key(key)
            if flag == '1':
                path = input(r'请输入需要加密的文件目录 '
                             r' C:\Users\Seven\Desktop\目录'
                             '\n')
                f = aes_file(path, key)
                print('文件加密成功！！！')
                time.sleep(3)
                q = input('输入任意键继续，输入 q 退出程序')
            elif flag == '2':
                try:
                    dpath = input(r'请输入需要解密的文件目录 '
                                  r' C:\Users\Seven\Desktop\目录'
                                  '\n')
                    d = decrypt_file(dpath, key)
                    print('文件解密成功！！！')
                    q = input('输入任意键继续，输入 q 退出程序')
                except UnicodeDecodeError:
                    print('密码不正确，请确认！！！')
                    time.sleep(1)
                    continue
            else:
                print('请输入正确的操作码')
                time.sleep(1)
                continue
    except FileNotFoundError:
        time.sleep(5)
        print('系统找不到指定的路径。')
    time.sleep(5)
    sys.exit()
